import featgen
import gengraph
import random
from torch_geometric.utils import from_networkx
from tqdm import tqdm
import numpy as np
import torch
from torch_geometric.utils.convert import to_networkx
from sklearn.model_selection import StratifiedKFold
import pdb

# 该函数用于计算数据集中的图数量。如果数据集中的图具有批次信息，则返回批次中的图数量；否则，返回特征张量的大小。
def num_graphs(data):
    if data.batch is not None:
        return data.num_graphs
    else:
        return data.x.size(0)

# 该函数用于将数据集划分为k折交叉验证的训练、验证和测试集。它使用StratifiedKFold进行分层划分，以保持类别之间的分布一致性。
def k_fold(dataset, folds, epoch_select):
    # 创建了一个StratifiedKFold对象，用于执行k折交叉验证。folds参数指定了将数据集分成多少份，shuffle=True表示在划分之前是否要对数据进行洗牌，random_state=12345设置了随机种子，以确保可复现性。
    skf = StratifiedKFold(folds, shuffle=True, random_state=12345)
    # 初始化两个空列表，用于存储测试集和训练集的索引。
    test_indices, train_indices = [], []
    # 使用skf.split()方法对数据集进行划分，其中torch.zeros(len(dataset))用于生成一个全零张量作为数据，dataset.data.y表示数据集的标签。skf.split()返回一个生成器，每次迭代返回两个数组的索引，分别表示训练集和测试集的索引。
    for _, idx in skf.split(torch.zeros(len(dataset)), dataset.data.y):
        # 将测试集的索引添加到test_indices列表中，并将NumPy数组转换为PyTorch张量。
        test_indices.append(torch.from_numpy(idx))
    # 根据epoch_select参数的值，确定如何选择验证集的索引。如果epoch_select为'test_max'，则将每个测试集的索引都作为验证集；否则，将每个测试集的前一个索引作为验证集。
    if epoch_select == 'test_max':
        val_indices = [test_indices[i] for i in range(folds)]
    else:
        val_indices = [test_indices[i - 1] for i in range(folds)]

    for i in range(folds):
        # 创建一个全1的张量作为训练集的掩码，数据类型为torch.uint8。
        train_mask = torch.ones(len(dataset), dtype=torch.uint8)
        # 将测试集的索引对应位置的值设置为0，以排除测试集中的样本。
        train_mask[test_indices[i].long()] = 0
        # 将验证集的索引对应位置的值设置为0，以排除验证集中的样本。
        train_mask[val_indices[i].long()] = 0
        # 将非零元素的索引添加到train_indices列表中，这将作为当前折的训练集索引。
        train_indices.append(train_mask.nonzero().view(-1))
    # 最后，返回了训练集、测试集和验证集的索引列表。
    return train_indices, test_indices, val_indices

# creat_one_pyg_graph() 函数的目的是创建一个PyTorch Geometric图形，并根据指定的参数生成特征。
# creat_one_pyg_graph() 函数接收多个参数，包括 context（图的类型，如 "tree" 或 "ba"）、shape（图的形状，如 "house"、"cycle" 等）、label（图的标签）、feature_dim（特征维度）、shape_num（图的数量）、settings_dict（设置字典，包含用于生成图的参数）和可选参数 args
def creat_one_pyg_graph(context, shape, label, feature_dim, shape_num, settings_dict, args=None):
    # 这里首先检查 args 是否为 None，如果是，则将 noise 设置为 0；否则，将 noise 设置为 args.noise。
    if args is None:
        noise = 0
    else:
        noise = args.noise
    # 如果 feature_dim 的值为 -1，则表示使用度作为特征。在这种情况下，创建了一个 ConstFeatureGen 实例，特征设置为 None，最大度设置为 args.max_degree。
    # 否则，如果 feature_dim 不为 -1，则表示使用随机生成的特征。在这种情况下，创建了一个 ConstFeatureGen 实例，特征是从均匀分布 [0, 1) 中随机生成的 feature_dim 维向量。
    if feature_dim == -1:
        # use degree as feature
        feature = featgen.ConstFeatureGen(None, max_degree=args.max_degree)
    else:
        feature = featgen.ConstFeatureGen(np.random.uniform(0, 1, feature_dim))
    # 接下来调用 gengraph.generate_graph() 函数来生成图。参数包括：
    # basis_type：图的类型，即 context。
    # shape：图的形状。
    # nb_shapes：要生成的图的数量。
    # width_basis：基础宽度，从 settings_dict 中获取。
    # feature_generator：特征生成器，即之前创建的 feature 对象。
    # m：参数 settings_dict 中的 m 值。
    # random_edges：随机边的数量，即之前计算的 noise 值。
    G, node_label = gengraph.generate_graph(basis_type=context,
                                            shape=shape,
                                            nb_shapes=shape_num,
                                            width_basis=settings_dict[context]["width_basis"],
                                            feature_generator=feature,
                                            m=settings_dict[context]["m"],
                                            random_edges=noise)
    # 使用 from_networkx() 函数将生成的 NetworkX 图转换为 PyTorch Geometric 图。
    # 为图设置标签 y，将其设置为一个包含当前图的标签的张量。
    pyg_G = from_networkx(G)
    pyg_G.y = torch.tensor([label])
    # 返回创建的 PyTorch Geometric 图 pyg_G 和节点标签 node_label。
    return pyg_G, node_label

# graph_dataset_generate() 函数的目的是生成用于训练的合成数据集，并将其保存到指定路径。
# 接收两个参数：args，包含数据集生成的参数；save_path，指定数据集保存的路径。
def graph_dataset_generate(args, save_path):
    # 定义了包含四种形状的类别列表 class_list，即 "house"、"cycle"、"grid" 和 "diamond"。
    # 定义了一个设置字典 settings_dict，其中包含两种图形 "ba" 和 "tree" 的相关参数设置。
    class_list = ["house", "cycle", "grid", "diamond"]
    settings_dict = {"ba": {"width_basis": args.node_num ** 2, "m": 2},
                     "tree": {"width_basis":2, "m": args.node_num}}
    # 提取参数 args 中的特征维度和图形数量，并初始化类别数量。
    # 创建一个空字典 dataset，其中包含键为 "tree" 和 "ba" 的子字典。
    feature_dim = args.feature_dim
    shape_num = args.shape_num
    class_num = class_list.__len__()
    dataset = {}
    dataset['tree'] = {}
    dataset['ba'] = {}
    # 对于类别列表中的每个形状，遍历其中的类别标签和形状名称，并输出当前正在创建的形状名称。
    for label, shape in enumerate(class_list):
        tr_list = []
        ba_list = []
        print("create shape:{}".format(shape))
        # 使用 tqdm 包装 range 函数，显示循环进度条。
        for i in tqdm(range(args.data_num)):
            # 对于每个形状，使用 creat_one_pyg_graph() 函数创建一个 "tree" 类型的图 tr_g 和一个 "ba" 类型的图 ba_g。
            tr_g, label1 = creat_one_pyg_graph(context="tree", shape=shape, label=label, feature_dim=feature_dim, 
                                               shape_num=shape_num, settings_dict=settings_dict, args=args)
            ba_g, label2 = creat_one_pyg_graph(context="ba", shape=shape, label=label, feature_dim=feature_dim, 
                                               shape_num=shape_num, settings_dict=settings_dict, args=args)
            # 将创建的 "tree" 和 "ba" 类型的图添加到相应的列表中。
            tr_list.append(tr_g)
            ba_list.append(ba_g)
        # 将包含 "tree" 和 "ba" 类型图的列表添加到 dataset 字典中，分别使用形状名称作为键。
        dataset['tree'][shape] = tr_list
        dataset['ba'][shape] = ba_list
    # 将数据集保存到指定路径，添加后缀名为 ".pt"。
    save_path += "/syn_dataset.pt"
    torch.save(dataset, save_path)
    # 输出保存路径信息。
    print("save at:{}".format(save_path))
    # 返回生成的合成数据集 dataset。
    return dataset

def test_dataset_generate(args, save_path):

    class_list = ["house", "cycle", "grid", "diamond"]
    settings_dict = {"ba": {"width_basis": (args.node_num) ** 2, "m": 2},
                     "tree": {"width_basis":2, "m": args.node_num}}

    feature_dim = args.feature_dim
    shape_num = args.shape_num
    class_num = class_list.__len__()
    dataset = {}
    dataset['tree'] = {}
    dataset['ba'] = {}
    data_num = int(0.2 * args.data_num)
    for label, shape in enumerate(class_list):
        tr_list = []
        ba_list = []
        print("test set create shape:{}".format(shape))
        for i in tqdm(range(data_num)):
            tr_g, label1 = creat_one_pyg_graph(context="tree", shape=shape, label=label, feature_dim=feature_dim, 
                                               shape_num=shape_num, settings_dict=settings_dict, args=args)
            ba_g, label2 = creat_one_pyg_graph(context="ba", shape=shape, label=label, feature_dim=feature_dim, 
                                               shape_num=shape_num, settings_dict=settings_dict, args=args)
            tr_list.append(tr_g)
            ba_list.append(ba_g)
        dataset['tree'][shape] = tr_list
        dataset['ba'][shape] = ba_list

    save_path += "/syn_dataset_test.pt"
    torch.save(dataset, save_path)
    print("save at:{}".format(save_path))
    return dataset

# dataset_bias_split() 函数的目的是根据给定的偏差参数和拆分比例将数据集划分为训练集、验证集和测试集，并确保每个类别的数据在各个数据集中的比例符合要求。
# 函数接受五个参数：dataset，原始数据集；args，包含数据集生成的参数；bias，用于指定数据集的偏差；split，指定训练集、验证集和测试集的拆分比例；total，指定总的样本数量。
def dataset_bias_split(dataset, args, bias=None, split=None, total=20000):
    # 定义了包含四种形状的类别列表 class_list，即 "house"、"cycle"、"grid" 和 "diamond"。
    # 创建了一个偏差字典 bias_dict，其中每个形状对应一个偏差值。
    class_list = ["house", "cycle", "grid", "diamond"]
    bias_dict = {"house": bias, "cycle": 1 - bias, "grid": 1 - bias, "diamond": 1 - bias}
    # 从原始数据集 dataset 中提取 "ba" 类型和 "tree" 类型的数据子集。
    ba_dataset = dataset['ba']
    tr_dataset = dataset['tree']
    # 将拆分比例转换为浮点数，并计算出训练集、验证集和测试集的样本数量。
    train_split, val_split, test_split = float(split[0]) / 10, float(split[1]) / 10, float(split[2]) / 10
    assert train_split + val_split + test_split == 1
    train_num, val_num, test_num = total * train_split, total * val_split, total * test_split
    # blance class
    # 从参数 args 中获取类别数量，并计算出每个类别在训练集、验证集和测试集中的样本数量。
    class_num = args.num_classes
    train_class_num, val_class_num, test_class_num = train_num / class_num, val_num / class_num, test_num / class_num
    # 创建空的训练集、验证集和测试集列表，并初始化边数为零。
    train_list, val_list, test_list  = [], [], []
    edges_num = 0
    # 对于每个形状，根据偏差值计算出每个类别在训练集、验证集和测试集中的样本数量。
    for shape in class_list:
        bias = bias_dict[shape]
        train_tr_num = int(train_class_num * bias)
        train_ba_num = int(train_class_num * (1 - bias))
        val_tr_num = int(val_class_num * bias)
        val_ba_num = int(val_class_num * (1 - bias))
        test_tr_num = int(test_class_num * 0.5)
        test_ba_num = int(test_class_num * 0.5)
        # 将每个类别的图添加到相应的训练集、验证集和测试集中。
        train_list += tr_dataset[shape][:train_tr_num] + ba_dataset[shape][:train_ba_num]
        val_list += tr_dataset[shape][train_tr_num:train_tr_num + val_tr_num] + ba_dataset[shape][train_ba_num:train_ba_num + val_ba_num]
        test_list += tr_dataset[shape][train_tr_num + val_tr_num:train_tr_num + val_tr_num + test_tr_num] + ba_dataset[shape][train_ba_num + val_ba_num:train_ba_num + val_ba_num + test_ba_num]
        # 计算每个类别图的边数，并将其添加到总边数中。
        _, e1 = print_graph_info(tr_dataset[shape][0], "Tree", shape)
        _, e2 = print_graph_info(ba_dataset[shape][0], "BA", shape)
        edges_num += e1 + e2
    # 对训练集、验证集和测试集进行洗牌，以确保数据的随机性。
    random.shuffle(train_list)
    random.shuffle(val_list)
    random.shuffle(test_list)
    # 计算每个类别图的平均边数，并返回划分好的训练集、验证集和测试集以及平均边数。
    the = float(edges_num) / (class_num * 2)
    return train_list, val_list, test_list, the

def print_graph_info(G, c, o):
    print('-' * 100)
    print("| graph: {}-{} | nodes num:{} | edges num:{} |".format(c, o, G.num_nodes, G.num_edges))
    print('-' * 100)
    return G.num_nodes, G.num_edges

def print_dataset_info(train_set, val_set, test_set, the):
    class_list = ["house", "cycle", "grid", "diamond"]
    dataset_group_dict = {}
    dataset_group_dict["Train"] = dataset_context_object_info(train_set, "Train", class_list, the)
    dataset_group_dict["Val"] = dataset_context_object_info(val_set, "Val   ", class_list, the)
    dataset_group_dict["Test"] = dataset_context_object_info(test_set, "Test  ", class_list, the)
    return dataset_group_dict

def dataset_context_object_info(dataset, title, class_list, the):

    class_num = len(class_list)
    tr_list = [0] * class_num
    ba_list = [0] * class_num
    for g in dataset:
        if g.num_edges > the: # ba
            ba_list[g.y.item()] += 1
        else: # tree
            tr_list[g.y.item()] += 1
    total = sum(tr_list) + sum(ba_list)
    info = "{} Total:{}\n| Tree: House:{:<5d}, Cycle:{:<5d}, Grids:{:<5d}, Diams:{:<5d} \n" +\
                        "| BA  : House:{:<5d}, Cycle:{:<5d}, Grids:{:<5d}, Diams:{:<5d} \n" +\
                        "| All : House:{:<5d}, Cycle:{:<5d}, Grids:{:<5d}, Diams:{:<5d} \n" +\
                        "| BIAS: House:{:.1f}%, Cycle:{:.1f}%, Grids:{:.1f}%, Diams:{:.1f}%"
    print("-" * 150)
    print(info.format(title, total, tr_list[0], tr_list[1], tr_list[2], tr_list[3],
                                    ba_list[0], ba_list[1], ba_list[2], ba_list[3],
                                    tr_list[0] +  ba_list[0],    
                                    tr_list[1] +  ba_list[1], 
                                    tr_list[2] +  ba_list[2], 
                                    tr_list[3] +  ba_list[3],
                                    100 *float(tr_list[0]) / (tr_list[0] +  ba_list[0]),
                                    100 *float(tr_list[1]) / (tr_list[1] +  ba_list[1]),
                                    100 *float(tr_list[2]) / (tr_list[2] +  ba_list[2]),
                                    100 *float(tr_list[3]) / (tr_list[3] +  ba_list[3]),
                     ))
    print("-" * 150)
    total_list = ba_list + tr_list
    group_counts = torch.tensor(total_list).float()
    return group_counts